from GNN import GNNClassifier as Classifier
import torch
import torch.nn as nn
from torch.utils.data import Dataset
from torch.utils.data import DataLoader
import os
import numpy as np
import random
# ======modify config:
CLASS_NUM = 6
# ===============================================data
# utils:


def get_one_hot(id, total_num):
    one_hot = np.zeros(total_num)
    one_hot[id] = 1
    return one_hot

# dataset


class nodeDataset(Dataset):
    def __init__(self, root_path):
        # load data and labels
        class_names = os.listdir(root_path)
        class_unique = list(set(class_names))
        class_num = len(class_unique)
        class_to_id = {name: i for i, name in enumerate(class_unique)}
        self.data_ = []
        self.label_ = []
        for name in class_names:
            all_data = np.load(root_path+"/"+name)
            for d in all_data:
                self.data_.append(d)
                self.label_.append(class_to_id[name])
        rand_index = list(range(len(self.data_)))
        random.shuffle(rand_index)
        self.data = []
        self.label = []
        self.data_label = []
        for i in range(len(self.data_)):
            self.data.append(self.data_[rand_index[i]])
            self.label.append(self.label_[rand_index[i]])
            self.data_label.append((self.data[-1], self.label[-1]))

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        return self.data_label[idx]


dataset = nodeDataset("./data/nodes")

# dataloader:
dataLoader = DataLoader(dataset, 64, shuffle=True)

# A mtrix:
A = torch.from_numpy(np.load('./data/A.npy'))
# ==============================================model
# load model
model = Classifier(21, 3, 512, 64, CLASS_NUM)
loss_fn = nn.CrossEntropyLoss()
# ==============================================train
# config:
EPOCH = 200
opt = torch.optim.Adam(model.parameters(), lr=0.01)

# start train：
model.train()
for epoch in range(EPOCH):
    ep_loss = 0
    for x, y in dataLoader:
        opt.zero_grad()
        y_ = model(A, x)
        y = y.to(dtype=torch.long)
        loss = loss_fn(y_, y)
        loss.backward()
        opt.step()
        ep_loss += loss.item()
    print("EPOCH: ", epoch, " loss: ", ep_loss)

torch.save(model, './model.npt')
